<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>第四章：Chains模块 - 4.4 Router Chains：根据输入动态选择下一个Chain (进阶)</title>
    <link rel="stylesheet" href="../shared_styles.css">
</head>
<body>
    <div class="container">
        <h2>4.4 Router Chains：根据输入动态选择下一个Chain (进阶)</h2>
            <p>Router Chain（路由链）是一种强大的机制，它允许我们根据用户的输入内容或预设的条件，动态地决定将输入路由到哪个“子链”（Sub-chain）或“目标”（Destination）进行处理。这使得构建能够处理多种类型输入、并根据输入特性采取不同行动的复杂应用成为可能。</p>
            <p>在Langchain的早期版本中，<code>MultiPromptChain</code> 和 <code>LLMRouterChain</code> 提供了路由功能。随着Langchain Expression Language (LCEL) 的引入，构建路由链变得更加灵活和直观，通常通过结合使用 <code>RunnableBranch</code> 和一个负责路由决策的逻辑来实现。</p>
            <p>本节我们将深入探讨如何使用LCEL和<code>RunnableBranch</code>来构建一个根据用户问题类型（物理、数学、历史等）选择不同处理逻辑的路由示例。我们将看到如何定义不同的处理分支，如何构建路由决策逻辑，以及如何将它们组合成一个完整的路由链。</p>
            <pre><code class="language-python">
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain.chains.router.llm_router import RouterOutputParser # 解析路由LLM的输出
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE # 预置的路由提示模板
from langchain_core.runnables import RunnableBranch, RunnableLambda, RunnablePassthrough
from langchain_core.messages import AIMessage # 用于类型提示
from dotenv import load_dotenv
import os

# (Qwen LLM 初始化代码同上, temperature=0 用于路由判断)
load_dotenv()
api_key = os.environ.get("DASHSCOPE_API_KEY")
qwen_llm_router = ChatOpenAI(base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", api_key=api_key, model="qwen-plus", temperature=0)
qwen_llm_qa = ChatOpenAI(base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", api_key=api_key, model="qwen-plus", temperature=0.7) # 用于问答

# 1. 定义不同场景的提示模板和对应的Runnable (prompt | llm)
physics_template_str = """你是一位物理学专家。用户用通俗易懂的语言问了关于物理的问题。
请用简洁明了、易于理解的方式回答。如果问题超出物理范畴，请说明。
问题：
{input}
回答："""
physics_prompt = PromptTemplate.from_template(physics_template_str)
physics_runnable = physics_prompt | qwen_llm_qa

math_template_str = """你是一位数学家。用户问了一个数学问题。
请清晰地解释概念并给出答案。如果问题非常复杂，可以建议用户分解问题。
问题：
{input}
回答："""
math_prompt = PromptTemplate.from_template(math_template_str)
math_runnable = math_prompt | qwen_llm_qa

history_template_str = """你是一位历史爱好者。用户对某个历史事件或时期感兴趣。
请提供相关的历史背景和关键信息。
问题：
{input}
回答："""
history_prompt = PromptTemplate.from_template(history_template_str)
history_runnable = history_prompt | qwen_llm_qa

# 目标信息，用于生成路由提示
prompt_infos_for_router = [
    {
        "name": "物理问答",
        "description": "适用于回答关于物理学原理、现象或概念的问题",
    },
    {
        "name": "数学问答",
        "description": "适用于回答数学问题，如代数、几何、微积分等",
    },
    {
        "name": "历史问答",
        "description": "适用于回答关于历史事件、人物或时期的问题",
    },
]

# 2. 创建路由逻辑 (一个Runnable，它会决定将用户输入路由到哪个目标)
destinations_for_router_prompt = [f"{p['name']}: {p['description']}" for p in prompt_infos_for_router]
destinations_str_for_router = "\n".join(destinations_for_router_prompt)

router_template_str_filled = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str_for_router)
router_prompt_lcel = PromptTemplate(
    template=router_template_str_filled,
    input_variables=["input"],
    output_parser=RouterOutputParser(), 
)

# 路由决策 Runnable: 输入 {"input": "用户问题"} -> 输出 {"destination": "名称", "next_inputs": {"input":"用户问题"}}
router_decision_runnable = router_prompt_lcel | qwen_llm_router | RouterOutputParser()

# 3. 定义RunnableBranch的条件和分支
# RunnableBranch的核心在于根据条件将执行流导向不同的Runnable。
# 每个分支由一个条件（condition）和一个Runnable（branch_runnable）组成。
# 当条件满足时，对应的branch_runnable就会被执行。

# 条件函数：这些函数接收RunnableBranch的输入，并返回一个布尔值。
# 在我们的例子中，RunnableBranch的输入是router_decision_runnable的输出字典
# (即 {"destination": "...", "next_inputs": {"input": "..."}})。
# 我们的条件需要检查这个字典中的 "destination" 键是否与某个分支的名称匹配。
def check_destination(name: str):
    # RunnableLambda 允许我们将一个简单的函数转换为一个Runnable。
    # 这里的lambda函数接收router_decision_runnable的输出 x，并检查 x["destination"] 是否等于 name。
    return RunnableLambda(lambda x: x.get("destination") == name)

# 分支的输入处理：每个分支的Runnable需要接收正确的输入。
# router_decision_runnable 的输出是一个包含 "next_inputs" 键的字典，
# 这个 "next_inputs" 字典才是我们希望传递给下游分支Runnable的实际输入。
def get_next_inputs(router_output: dict):
    # 这个函数从路由决策的输出字典中提取 "next_inputs" 字典。
    # 如果 "next_inputs" 不存在（尽管在RouterOutputParser的输出中应该存在），则提供一个默认值。
    return router_output.get("next_inputs", {"input": ""}) # Fallback if next_inputs is missing


# 默认的Runnable：如果所有带条件的branches都不匹配，RunnableBranch会执行最后一个没有条件的Runnable。
# 这里我们定义一个简单的回退Runnable，使用默认的PromptTemplate和LLM来尝试回答问题。
default_fallback_prompt = PromptTemplate.from_template("请尝试回答用户的问题：{input}")
default_fallback_runnable = default_fallback_prompt | qwen_llm_qa

# RunnableBranch的构建：
# RunnableBranch 接收一系列 (condition, branch_runnable) 对，以及一个可选的默认Runnable。
# 它会按顺序评估每个条件，执行第一个条件为 True 的 branch_runnable。
# 如果所有条件都为 False，则执行默认Runnable。
# 注意：每个 branch_runnable 的输入是 RunnableBranch 接收到的输入。
# 在我们的例子中，RunnableBranch 接收的是 router_decision_runnable 的输出字典。
# 因此，我们需要在每个分支的 branch_runnable 前面加上 RunnableLambda(get_next_inputs)
# 来从这个字典中提取出真正的输入 {"input": "..."} 传递给下游的 physics_runnable, math_runnable 等。
branch = RunnableBranch(
    # 如果路由决策的 destination 是 "物理问答"，则执行这个分支。
    # RunnableLambda(get_next_inputs) 提取 {"input": "..."}，然后传递给 physics_runnable。
    (check_destination("物理问答"), RunnableLambda(get_next_inputs) | physics_runnable),
    # 如果路由决策的 destination 是 "数学问答"，则执行这个分支。
    (check_destination("数学问答"), RunnableLambda(get_next_inputs) | math_runnable),
    # 如果路由决策的 destination 是 "历史问答"，则执行这个分支。
    (check_destination("历史问答"), RunnableLambda(get_next_inputs) | history_runnable),
    # 如果以上条件都不满足，执行默认回退分支。
    # 同样，需要提取 {"input": "..."} 传递给 default_fallback_runnable。
    RunnableLambda(get_next_inputs) | default_fallback_runnable
)

# 4. 构建完整的路由链 (LCEL)
# 完整的路由链需要完成以下步骤：
# a. 接收原始用户输入 {"input": "用户问题"}。
# b. 将用户输入传递给路由决策Runnable (router_decision_runnable)。
# c. 将路由决策Runnable的输出传递给RunnableBranch。
# d. RunnableBranch 根据决策输出选择并执行相应的分支。

# 使用 RunnablePassthrough.assign 来在链中保留原始输入，并添加路由决策的输出。
# RunnablePassthrough.assign({ new_key: runnable }) 会执行 runnable，
# 并将 runnable 的输出赋值给 new_key，同时保留原始输入。
# 这里的 runnable 是 router_decision_runnable，它接收原始输入 {"input": "..."} 并输出路由决策字典。
# 结果是一个包含原始输入和路由决策输出的字典，例如：
# {"input": "用户问题", "router_output": {"destination": "...", "next_inputs": {"input": "..."}}}
full_router_lcel = RunnablePassthrough.assign(
    router_output=router_decision_runnable
) | \
# 接下来，我们需要将 router_output (即路由决策字典) 传递给 RunnableBranch。
# 使用 RunnableLambda 来从上一步的字典中提取 "router_output" 键的值。
RunnableLambda(lambda x: x["router_output"]) | \
# 最后，将提取出的路由决策字典传递给 branch。
branch


# 5. 测试Router Chain
def test_router(user_input):
    print(f"用户输入: {user_input}")
    # full_router_lcel 期望一个字典作为输入，其中包含 "input" 键。
    response_message = full_router_lcel.invoke({"input": user_input})
    print(f"LLM回答: {response_message.content}")
    print("\n" + "="*30 + "\n")

test_router("什么是黑洞？")
test_router("求解二次方程 x^2 - 5x + 6 = 0")
test_router("法国大革命是什么时候开始的？")
test_router("今天天气怎么样？") # 应该会走到default_fallback_runnable
            </code></pre>
            <p>通过上述LCEL的组合，我们成功构建了一个Router Chain。整个流程如下：</p>
            <ol>
                <li>用户输入通过 <code>RunnablePassthrough.assign</code> 传递给 <code>router_decision_runnable</code>。</li>
                <li><code>router_decision_runnable</code> 根据用户输入和预设的路由提示，决定将输入路由到哪个目标，并输出一个包含目标名称和传递给下一阶段输入的字典。</li>
                <li><code>RunnableLambda(lambda x: x["router_output"])</code> 从上一步的输出中提取路由决策字典。</li>
                <li>这个路由决策字典被传递给 <code>RunnableBranch</code>。</li>
                <li><code>RunnableBranch</code> 评估其内部定义的条件（使用 <code>check_destination</code> 函数），找到第一个匹配的条件。</li>
                <li>执行与匹配条件对应的分支Runnable。每个分支的Runnable都以 <code>RunnableLambda(get_next_inputs)</code> 开始，确保从路由决策字典中提取出正确的输入传递给下游的Prompt和LLM。</li>
                <li>如果没有任何条件匹配，则执行默认的回退分支。</li>
            </ol>
            <p>为了更好地理解Router Chain的工作流程，下面是一个简化的流程图：</p>
            <pre><code class="language-mermaid">
graph TD
    A[用户输入] --> B{路由决策};
    B -->|物理问题| C[物理问答Runnable];
    B -->|数学问题| D[数学问答Runnable];
    B -->|历史问题| E[历史问答Runnable];
    B -->|其他问题| F[默认回退Runnable];
    C --> G[LLM回答];
    D --> G;
    E --> G;
    F --> G;
    G --> H[最终输出];
</code></pre>
            <p>这种基于LCEL和<code>RunnableBranch</code>的路由方式非常灵活，可以轻松扩展以支持更多的路由目标和更复杂的路由逻辑。</p>
        <div class="navigation">
            <a href="../langchain_tutorial_outline.html">返回目录</a>
            <a href="../chapter4/4-3_sequential_chains.html">上一小节：Sequential Chains：按顺序链接多个Chain</a>
            <a href="../chapter4/4-5_custom_chain.html">下一小节：自定义Chain</a>
        </div>
    </div>
</body>
</html>
